#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
只保留一个实例脚本
确保GPU服务器至少保持一个实例运行
无论当前实例数量是多少（0个或大于1个），都会恢复到1个实例
"""

import requests
import json
import time
import sys
import argparse
from typing import List, Dict, Any, Optional

class KeepOneInstanceManager:
    """只保留一个实例管理器"""
    
    def __init__(self, base_url: str = "http://localhost:8080"):
        self.base_url = base_url
    
    def get_servers(self) -> List[Dict[str, Any]]:
        """获取所有服务器列表"""
        try:
            response = requests.get(f"{self.base_url}/api/v1/servers", timeout=10)
            if response.status_code == 200:
                data = response.json()
                if data.get('success'):
                    return data.get('data', [])
            else:
                print(f"❌ 获取服务器列表失败: {response.status_code}")
                return []
        except Exception as e:
            print(f"❌ 连接失败: {e}")
            return []
    
    def create_server(self, region: str = "ap-beijing", gpu_type: str = "basic") -> Optional[str]:
        """创建新服务器"""
        try:
            print(f"🚀 创建新服务器: {region} - {gpu_type}")
            response = requests.post(f"{self.base_url}/api/v1/servers", 
                                   json={'region': region, 'gpu_type': gpu_type}, 
                                   timeout=60)
            if response.status_code == 200:
                data = response.json()
                if data.get('success'):
                    instance_id = data.get('instance_id')
                    print(f"✅ 服务器创建成功: {instance_id}")
                    return instance_id
                else:
                    print(f"❌ 服务器创建失败: {data.get('error')}")
                    return None
            else:
                print(f"❌ 创建请求失败: {response.status_code}")
                return None
        except Exception as e:
            print(f"❌ 创建异常: {e}")
            return None
    
    def destroy_server(self, instance_id: str) -> bool:
        """销毁服务器"""
        try:
            print(f"🗑️ 销毁服务器: {instance_id}")
            response = requests.delete(f"{self.base_url}/api/v1/servers/{instance_id}", timeout=30)
            if response.status_code == 200:
                data = response.json()
                if data.get('success'):
                    print(f"✅ 服务器销毁成功: {instance_id}")
                    return True
                else:
                    print(f"❌ 服务器销毁失败: {data.get('error')}")
                    return False
            else:
                print(f"❌ 销毁请求失败: {response.status_code}")
                return False
        except Exception as e:
            print(f"❌ 销毁异常: {e}")
            return False
    
    def get_server_status(self, instance_id: str) -> Optional[str]:
        """获取服务器状态"""
        try:
            response = requests.get(f"{self.base_url}/api/v1/servers/{instance_id}/status", timeout=10)
            if response.status_code == 200:
                data = response.json()
                if data.get('success'):
                    return data.get('data', {}).get('status')
            return None
        except:
            return None
    
    def wait_for_server_ready(self, instance_id: str, timeout: int = 300) -> bool:
        """等待服务器就绪"""
        print(f"⏳ 等待服务器就绪: {instance_id}")
        start_time = time.time()
        
        while time.time() - start_time < timeout:
            status = self.get_server_status(instance_id)
            if status == "RUNNING":
                print(f"✅ 服务器已就绪: {instance_id}")
                return True
            elif status in ["STOPPED", "TERMINATED"]:
                print(f"❌ 服务器状态异常: {instance_id} - {status}")
                return False
            
            print(f"⏳ 服务器状态: {status}，继续等待...")
            time.sleep(10)
        
        print(f"⏰ 等待超时: {instance_id}")
        return False
    
    def select_server_to_keep(self, servers: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
        """选择要保留的服务器"""
        if not servers:
            return None
        
        # 优先选择运行中的服务器
        running_servers = [s for s in servers if s.get('status') == 'RUNNING']
        if running_servers:
            # 选择最早创建的运行中服务器
            return min(running_servers, key=lambda x: x.get('created_time', ''))
        
        # 如果没有运行中的，选择状态最好的
        status_priority = ['PENDING', 'STOPPING', 'STOPPED', 'TERMINATED']
        for status in status_priority:
            status_servers = [s for s in servers if s.get('status') == status]
            if status_servers:
                return min(status_servers, key=lambda x: x.get('created_time', ''))
        
        # 如果都没有，返回第一个
        return servers[0]
    
    def keep_one_instance(self, preferred_region: str = "ap-beijing", preferred_gpu_type: str = "basic") -> bool:
        """只保留一个实例"""
        print("🎯 开始执行：只保留一个实例")
        print("=" * 50)
        
        # 获取当前服务器列表
        servers = self.get_servers()
        if servers is None:
            print("❌ 无法获取服务器列表")
            return False
        
        current_count = len(servers)
        print(f"📊 当前服务器数量: {current_count}")
        
        if current_count == 1:
            print("✅ 已经只有一个实例，无需调整")
            return True
        
        if current_count == 0:
            # 没有实例，创建一个
            print("🚀 没有实例，创建新实例...")
            instance_id = self.create_server(preferred_region, preferred_gpu_type)
            if instance_id:
                # 等待实例就绪
                if self.wait_for_server_ready(instance_id):
                    print("✅ 成功创建并启动一个实例")
                    return True
                else:
                    print("❌ 实例创建失败或未就绪")
                    return False
            else:
                print("❌ 实例创建失败")
                return False
        
        else:
            # 有多个实例，保留一个，销毁其他的
            print(f"🗑️ 有 {current_count} 个实例，需要保留1个，销毁 {current_count-1} 个")
            
            # 选择要保留的服务器
            server_to_keep = self.select_server_to_keep(servers)
            if not server_to_keep:
                print("❌ 无法选择要保留的服务器")
                return False
            
            keep_instance_id = server_to_keep.get('instance_id')
            print(f"✅ 选择保留服务器: {keep_instance_id}")
            
            # 销毁其他服务器
            destroyed_count = 0
            for server in servers:
                instance_id = server.get('instance_id')
                if instance_id != keep_instance_id:
                    if self.destroy_server(instance_id):
                        destroyed_count += 1
                    time.sleep(2)  # 避免并发销毁
            
            print(f"✅ 成功销毁 {destroyed_count} 个实例")
            
            # 等待系统稳定，多次检查
            print("⏳ 等待系统稳定...")
            max_retries = 12  # 最多等待120秒
            retry_count = 0
            
            while retry_count < max_retries:
                time.sleep(10)  # 每次等待10秒
                retry_count += 1
                
                final_servers = self.get_servers()
                if not final_servers:
                    print(f"📊 检查 {retry_count}/{max_retries}: 无法获取服务器列表")
                    continue
                
                final_count = len(final_servers)
                running_count = len([s for s in final_servers if s.get('status') == 'RUNNING'])
                pending_count = len([s for s in final_servers if s.get('status') == 'PENDING'])
                
                print(f"📊 检查 {retry_count}/{max_retries}: 总数 {final_count}, 运行中 {running_count}, 启动中 {pending_count}")
                
                # 如果只有一个运行中的实例，认为成功
                if running_count == 1 and final_count == 1:
                    print("✅ 成功保留一个实例")
                    return True
                elif running_count == 1 and final_count > 1:
                    # 有一个运行中的，但还有其他状态的实例，继续等待
                    print("⏳ 有一个运行中的实例，等待其他实例完全销毁...")
                    continue
                elif final_count == 0:
                    print("⚠️ 所有实例都被销毁了，需要创建一个新实例")
                    # 创建一个新实例
                    instance_id = self.create_server(preferred_region, preferred_gpu_type)
                    if instance_id and self.wait_for_server_ready(instance_id):
                        print("✅ 成功创建并启动一个实例")
                        return True
                    else:
                        print("❌ 创建新实例失败")
                        return False
            
            # 最终检查
            final_servers = self.get_servers()
            final_count = len(final_servers) if final_servers else 0
            print(f"📊 最终服务器数量: {final_count}")
            
            if final_count == 1:
                print("✅ 成功保留一个实例")
                return True
            else:
                print(f"⚠️ 保留实例未完全成功，当前有 {final_count} 个实例")
                return False
    
    def show_current_status(self):
        """显示当前状态"""
        print("📊 当前实例状态")
        print("=" * 50)
        
        servers = self.get_servers()
        if servers is None:
            print("❌ 无法获取服务器列表")
            return
        
        current_count = len(servers)
        print(f"📊 当前服务器数量: {current_count}")
        
        if current_count == 0:
            print("⚠️ 没有运行中的实例")
        elif current_count == 1:
            print("✅ 正好有一个实例")
        else:
            print(f"⚠️ 有 {current_count} 个实例，建议只保留一个")
        
        if servers:
            print("\n📋 实例列表:")
            for i, server in enumerate(servers, 1):
                instance_id = server.get('instance_id', 'Unknown')
                status = server.get('status', 'Unknown')
                region = server.get('region', 'Unknown')
                gpu_type = server.get('gpu_type', 'Unknown')
                print(f"  {i:2d}. {instance_id} | {status} | {region} | {gpu_type}")

def main():
    parser = argparse.ArgumentParser(description='只保留一个实例脚本')
    parser.add_argument('action', choices=['keep-one', 'status'], 
                       help='操作类型: keep-one(保留一个实例), status(查看状态)')
    parser.add_argument('--region', default='ap-beijing', 
                       help='首选地域 (默认: ap-beijing)')
    parser.add_argument('--gpu-type', default='basic', 
                       help='首选GPU类型 (默认: basic)')
    parser.add_argument('--url', default='http://localhost:8080', 
                       help='HAI优化系统URL')
    
    args = parser.parse_args()
    
    manager = KeepOneInstanceManager(args.url)
    
    if args.action == 'keep-one':
        success = manager.keep_one_instance(args.region, args.gpu_type)
        if success:
            print("\n🎉 操作完成：成功保留一个实例")
            sys.exit(0)
        else:
            print("\n❌ 操作失败：无法保留一个实例")
            sys.exit(1)
    elif args.action == 'status':
        manager.show_current_status()

if __name__ == '__main__':
    main()
